package huanzhao.erp.other.aspect;

import cn.hutool.core.net.Ipv4Util;
import huanzhao.erp.domain.exception.RateLimitException;
import huanzhao.erp.domain.po.Access;
import huanzhao.erp.domain.annotation.RateLimit;
import net.jodah.expiringmap.ExpirationPolicy;
import net.jodah.expiringmap.ExpiringMap;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.aop.aspectj.MethodInvocationProceedingJoinPoint;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

/**
 * 切面类：实现限流校验
 * 用于需要限流的接口较少的情况
 * 如果需要限流的接口较多，建议使用拦截器限流
 *
 * @author April
 * @version 1.0
 * 2024/7/25 下午1:06
 */
@Aspect
@Component
public class AccessLimitAspect {
    public static final ExpiringMap<String, Access> IP_ACCESS_CACHE = ExpiringMap.builder().expirationPolicy(ExpirationPolicy.CREATED).variableExpiration().build();
    private static final String[] proxyHeaders = {"x-forwarded-for", "Proxy-Client-IP", "WL-Proxy-Client-IP"};
    private static final Logger log = LoggerFactory.getLogger(AccessLimitAspect.class);

    /**
     * 定义切点
     */
    @Pointcut("@annotation(huanzhao.erp.domain.annotation.RateLimit)")
    public void limitPointCut() {
    }

    /**
     * 环绕通知
     */
    @Around("limitPointCut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        // 获取被注解的方法
        MethodInvocationProceedingJoinPoint mjp = (MethodInvocationProceedingJoinPoint) joinPoint;
        MethodSignature signature = (MethodSignature) mjp.getSignature();
        Method method = signature.getMethod();
        // 获取方法上的注解
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        // 获取request对象
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes == null) {
            return joinPoint.proceed();
        }
        HttpServletRequest request = attributes.getRequest();
        String key = "rl:ip:" + getRawIpAddress(request) + ":" + request.getRequestURI().hashCode();
        Access access = IP_ACCESS_CACHE.get(key);
        if (access != null) {
            long ttl = rateLimit.period() + (access.getCreateTime() - System.currentTimeMillis()) / 1000;
            // 超出访问次数限制
            if (access.getCount() >= rateLimit.rate()) {
                // 获取key的ttl
                throw new RateLimitException(rateLimit.msg(), ttl);
            }
            access.setCount(access.getCount() + 1);
            // 访问次数自增1
            IP_ACCESS_CACHE.put(key, access, ExpirationPolicy.CREATED, ttl, TimeUnit.SECONDS);
        } else {
            // 第一次设置数据,过期时间为注解确定的访问周期
            IP_ACCESS_CACHE.put(key, new Access(System.currentTimeMillis(), 1), ExpirationPolicy.CREATED, rateLimit.period(), TimeUnit.SECONDS);
        }
        return joinPoint.proceed();
    }

    /**
     * 获取真实的、未经过代理的IP地址
     *
     * @param request 请求对象
     * @return long ip地址对应的整数值
     * 2024/7/25 下午1:18:29
     */
    private long getRawIpAddress(HttpServletRequest request) {
        String ipAddress = null;
        boolean isProxied = false;
        for (String header : proxyHeaders) {
            ipAddress = request.getHeader(header);
            // ip地址长度大于等于15，说明ip可能有效
            if (ipAddress != null && ipAddress.length() >= 15) {
                isProxied = true;
                break;
            }
        }
        // 如果没有代理，直接返回远程IP地址
        if (!isProxied) {
            return Ipv4Util.ipv4ToLong(request.getRemoteAddr());
        }
        // 如果ip地址长度大于15且包含逗号，说明有多个代理，取第一个ip地址
        if (ipAddress.length() > 15 && ipAddress.indexOf(",") > 0) {
            ipAddress = ipAddress.substring(0, ipAddress.indexOf(","));
        }
        try {
            return Ipv4Util.ipv4ToLong(ipAddress);
        } catch (IllegalArgumentException e) {
            log.warn("转换IP地址失败，待转换的IP地址为：{}，远程IP地址为：{}", ipAddress, request.getRemoteAddr());
            return Ipv4Util.ipv4ToLong(request.getRemoteAddr());
        }
    }
}
